import numpy as np
import cv2
import onnxruntime as ort
import matplotlib.pyplot as plt
import torch
from yoloseg import YOLOSeg


if __name__ == '__main__':

    #model = ort.InferenceSession("wheel_seg_best.onnx",  providers=['CUDAExecutionProvider'])
    image = plt.imread("Data/50.bmp")
    yoloseg=YOLOSeg("wheel_seg_best.onnx")
    input_tensor=yoloseg.prepare_input(image)
    outputs = yoloseg.inference(input_tensor)
    yoloseg.boxes, yoloseg.scores, yoloseg.class_ids, mask_pred=yoloseg.process_box_output(outputs[0])
    mask_maps=yoloseg.process_mask_output(mask_pred, outputs[1])

    inner_mask=np.uint8(mask_maps[yoloseg.class_ids==1].squeeze()*255)
    outter_mask=np.uint8(mask_maps[yoloseg.class_ids==2].squeeze()*255)

    # plt.imshow(inner_mask,cmap="gray")
    # plt.show()
    #
    # plt.imshow(outter_mask,cmap="gray")
    # plt.show()


    contours, hierarchy = cv2.findContours(inner_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours_num_list = []
    for i in range(len(contours)):
        contours_num_list.append(contours[i].shape[0])
    ellipse = cv2.fitEllipse(contours[np.argmax(np.array(contours_num_list))])
    # cv2.ellipse(img_source, ellipse, (0, 0, 255), 1, cv2.LINE_AA)
    x, y = ellipse[0]
    r1 = np.mean(ellipse[1]) / 2*1.01
    cv2.circle(image, (int(x), int(y)), int(r1), (255, 0, 0), 5)

    contours, hierarchy = cv2.findContours(outter_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours_num_list = []
    for i in range(len(contours)):
        contours_num_list.append(contours[i].shape[0])
    ellipse = cv2.fitEllipse(contours[np.argmax(np.array(contours_num_list))])
    # cv2.ellipse(img_source, ellipse, (0, 0, 255), 1, cv2.LINE_AA)
    x, y = ellipse[0]
    r2 = np.mean(ellipse[1]) / 2*1.01
    cv2.circle(image, (int(x), int(y)), int(r2), (0, 0, 255), 5)

    print("内轮半径为："+str(int(r1)),"; 外轮半径为："+str(int(r2)))
    plt.imshow(image)
    plt.show()

    # image=plt.imread("1.bmp")/255
    # height, width, _ = image.shape
    # length = max(height, width)
    # imageB = np.zeros((640, 640, 3), np.float32)
    # s = 640 / length
    #
    # resized_img = cv2.resize(image, (0, 0), fx=s, fy=s)
    # h, w, c = resized_img.shape
    # imageB[0: h, 0: w] = resized_img
    #
    # imageB = np.array([imageB.transpose(2, 0, 1)])
    #
    # output = model.run(None, {"images": imageB})
    # for i in range(32):
    #     plt.imshow(output[1][0,i,:,:],cmap="gray")
    #     plt.show()
